#include <pthread.h>
#include <errno.h>
#include <stdio.h>
#include <netdb.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>
#include <stdio.h>

#include "psd_conn.h"
#include "psd_protocols.h"
#include "psd_conn_udt.h"
#include "psd_logger.h"
#include "psd_tpool.h"
#include "psd_event.h"
#include "psd_config.h"
#include "psd_modules.h"

#include "udt_c_wrapper.h"

static psdConn_udtData *psd_conn_udtdata_alloc(UDTSOCKET sd);
static int psd_conn_udt_write(psdConn *sink, const void *buf, size_t len, int flags);
static int psd_conn_udt_read(psdConn *src, void *buf, size_t len, int flags);
static int psd_conn_udt_shutdown(psdConn *conn, uint8_t direction);
static int psd_conn_udt_setbufsize(psdConn *conn, uint8_t direction, int size);
static int psd_conn_udt_settimeout(psdConn *conn, uint8_t direction, int seconds);
static void psd_conn_udt_free(psdConn *conn);


psdConn *psd_conn_udt_alloc(UDTSOCKET sd) {
	psdConn *new_conn;
	char buf[200];

	new_conn = psd_conn_alloc();
	if (!new_conn)
		goto error_exit;

	new_conn->conn_private = psd_conn_udtdata_alloc(sd);
	if (!new_conn->conn_private)
		goto error_exit2;

	new_conn->stats_private = psd_conn_stats_alloc_def();
	if (!new_conn->stats_private) {
		psd_err(0, "couldn't allocate tcp stats structure");
		goto error_exit3;
	}

	new_conn->protocol = "UDT";

	sprintf(buf, "%d", sd);
	new_conn->description = strdup(buf);

	new_conn->read2 = psd_conn_udt_read;
	new_conn->write2 = psd_conn_udt_write;
	new_conn->shutdown2 = psd_conn_udt_shutdown;
	new_conn->setbufsize2 = psd_conn_udt_setbufsize;
	new_conn->settimeout2 = psd_conn_udt_settimeout;
	new_conn->get_stat2 = psd_conn_default_get_stat;
	new_conn->free_conn_private2 = psd_conn_udt_free;
	new_conn->free_stats2 = psd_conn_free_stats_def;
	new_conn->send_msg2 = psd_conn_default_send_msg;
	new_conn->get_msg2 = psd_conn_default_get_msg;
	new_conn->set_session_status2=psd_conn_default_set_session_status;
	return new_conn;

error_exit4:
	free(new_conn->stats_private);
error_exit3:
	free(new_conn->conn_private);
error_exit2:
	free(new_conn);
error_exit:
	return NULL;
}

psdConn_udtData *psd_conn_udtdata_alloc(UDTSOCKET sd) {
	psdConn_udtData *new_udtdata;

	new_udtdata = (psdConn_udtData *) malloc(sizeof(psdConn_udtData));
	if (!new_udtdata)
		goto error_exit;

	if (pthread_mutex_init(&(new_udtdata->lock), 0) < 0) {
		goto error_exit_udtdata;
	}

	new_udtdata->hdr_bufs = malloc(sizeof(psdUdtHeader) * 1000);
	if (!new_udtdata->hdr_bufs)
		goto error_exit_mutex;

	new_udtdata->sd = sd;
	new_udtdata->closed = 0;
	new_udtdata->read_closed = 0;
	new_udtdata->write_closed = 0;
	new_udtdata->left_to_read = 0;
	new_udtdata->hdr_bufs_head = 0;

	return new_udtdata;

error_exit_mutex:
	pthread_mutex_destroy(&new_udtdata->lock);
error_exit_udtdata:
	free(new_udtdata);
error_exit:
	return NULL;
}


static int psd_conn_udt_write(psdConn *sink, const void *buf, size_t len, int flags) {
	psdConn_udtData *udt_data = (psdConn_udtData *) sink->conn_private;
	int sent;

	psd_info(10, "udt_write()");

	pthread_mutex_lock(&(udt_data->lock));
	{
		psdUdtHeader *hdr = &(udt_data->hdr_bufs[udt_data->hdr_bufs_head]);

		udt_data->hdr_bufs_head++;
		udt_data->hdr_bufs_head %= 950;

		hdr->type = UDT_TYPE_DATA;
		hdr->args = htonl(len);

		if (udt_data->write_closed) {
			sent = -1;
		} else {
			psd_info(10, "Writing header");
			uint32_t amt_to_send = sizeof(psdUdtHeader);
			uint32_t amt_sent = 0;
			do {
				sent = udt_send(udt_data->sd, ((char *) hdr) + amt_sent, amt_to_send, flags, NULL, NULL, NULL);
				psd_info(10, "Sent: %d", sent);
				if (sent > 0) {
					amt_sent += sent;
					amt_to_send -= sent;
					psd_info(10, "Amt_sent: %d", amt_sent);
				} else {
					psd_info(10, "Failed to write header: %d", sent);
					break;
				}
			} while (amt_to_send > 0);

			psd_info(10, "Amt_sent: %d", amt_sent);
			if (amt_sent == sizeof(psdUdtHeader)) {
				psd_info(10, "Wrote header. Writing body: %ld", len);
				((psdConn_defStats *)sink->stats_private)->bytes_written += sent;
				sent = udt_send(udt_data->sd, (char *) buf, len, flags, NULL, NULL, NULL);
				if (sent > 0)
					((psdConn_defStats *)sink->stats_private)->bytes_written += sent;
			}
		}
	}
	pthread_mutex_unlock(&(udt_data->lock));

	return sent;
}

static int psd_conn_udt_read(psdConn *src, void *buf, size_t len, int flags) {
	psdConn_udtData *udt_data = (psdConn_udtData *) src->conn_private;
	int recvd;
	uint32_t amt_to_read;
	uint32_t amt_read;
	psdUdtHeader hdr;

	psd_info(10, "udt_read()");

	if (udt_data->read_closed == 1) {
		psd_err(0, "Tried to read from closed UDT socket");
		return -1;
	}

	while (udt_data->left_to_read == 0) {
		amt_to_read = sizeof(hdr);
		amt_read = 0;

		do {
			psd_info(10, "udt_recv(): in: %d", amt_to_read);
			recvd = udt_recv(udt_data->sd, ((char *) &hdr) + amt_read, amt_to_read, flags, NULL, NULL, NULL);
			psd_info(10, "udt_recv(): out: %d", recvd);
			if (recvd <= 0) {
				psd_err(0, "Failed to receive header: %d", recvd);
				return recvd;
			}

			amt_to_read -= recvd;
			amt_read += recvd;
		} while (amt_to_read > 0);

		if (hdr.type == UDT_TYPE_CLOSE) {
			uint32_t direction = ntohl(hdr.args);

			psd_info(10, "Got close message");

			if (direction == PSD_RECV_SIDE) {
				pthread_mutex_lock(&(udt_data->lock));
				{
					udt_data->write_closed = 1;
				}
				pthread_mutex_unlock(&(udt_data->lock));
			} else {
				udt_data->read_closed = 1;
				return 0;
			}
		} else if (hdr.type == UDT_TYPE_DATA) {
			udt_data->left_to_read = ntohl(hdr.args);
			psd_info(10, "Got data message for %lu bytes", udt_data->left_to_read);

		} else {
			psd_info(10, "got a junk type: %u", hdr.type);
			return -1;
		}
	}

	amt_to_read = udt_data->left_to_read;
	if (amt_to_read > len) {
		amt_to_read = len;
	}

	psd_info(10, "udt_recv(): in: %d", amt_to_read);
	recvd = udt_recv(udt_data->sd, buf, amt_to_read, flags, NULL, NULL, NULL);
	psd_info(10, "udt_recv(): out: %d", recvd);
	if (recvd > 0) {
		udt_data->left_to_read -= recvd;
		((psdConn_defStats *)src->stats_private)->bytes_read += recvd;
	}

	psd_info(10, "udt recv returned %d", recvd);

	return recvd;
}

static int psd_conn_udt_shutdown(psdConn *conn, uint8_t direction) {
	psdConn_udtData *udt_data = (psdConn_udtData *) conn->conn_private;

#if 1
	if (direction == PSD_SEND_SIDE || direction == PSD_RECV_SIDE) {
		psd_warn(0, "udt doesn't implement shutdown, only close");
	}

	pthread_mutex_lock(&(udt_data->lock));
	{
		udt_data->closed |= direction;
		if (direction == (PSD_SEND_SIDE | PSD_RECV_SIDE)) {
			gettimeofday(&(conn->end_time), NULL);
			psd_event("phoebus.connection.end", conn, "DEPOT=%s PROTOCOL=%s", conn->description, conn->protocol);
			udt_close(udt_data->sd);
		} else {
			psdUdtHeader *hdr = &(udt_data->hdr_bufs[udt_data->hdr_bufs_head++]);
			udt_data->hdr_bufs_head %= 950;

			hdr->type = UDT_TYPE_CLOSE;
			hdr->args = htonl((uint32_t) direction);

			udt_send(udt_data->sd, (char *) hdr, sizeof(psdUdtHeader), 0, NULL, NULL, NULL);
		}
	}
	pthread_mutex_unlock(&(udt_data->lock));

	return 0;
#else
	if (direction == PSD_SEND_SIDE) {
		udt_shutdown(udt_data->sd, SHUT_WR);
	} else if (direction == PSD_RECV_SIDE) {
		udt_shutdown(udt_data->sd, SHUT_RD);
	} else if (direction == (PSD_RECV_SIDE | PSD_SEND_SIDE)) {
		udt_close(udt_data->sd);
	} else {
		return -1;
	}
#endif

	return 0;
}

static int psd_conn_udt_setbufsize(psdConn *conn, uint8_t direction, int size) {
	psdConn_udtData *udt_data = (psdConn_udtData *) conn->conn_private;
	int new_bufsize;
	int junk;

	if (direction & PSD_RECV_SIDE) {
		if ((udt_setsockopt(udt_data->sd, SOL_SOCKET, UDT_RCVBUF, (const void *)&size, sizeof(int))) < 0) {
			psd_err(10, "set recv buf failed");
			goto error_exit;
		}

		if ((udt_getsockopt(udt_data->sd, SOL_SOCKET, UDT_RCVBUF, (char *)&new_bufsize, &junk)) == 0) {
			psd_info(10, "recv buf set to \"%d\"", new_bufsize);
		}
	} 
	
	if (direction & PSD_SEND_SIDE) {
		if ((udt_setsockopt(udt_data->sd, SOL_SOCKET, UDT_SNDBUF, (const void *)&size, sizeof(int))) < 0) {
			psd_err(10, "set recv buf failed");
			goto error_exit;
		}

		if ((udt_getsockopt(udt_data->sd, SOL_SOCKET, UDT_SNDBUF, (char *)&new_bufsize, &junk)) == 0) {
			psd_info(10, "send buf set to \"%d\"", new_bufsize);
		}
	} 

	return 0;

error_exit:
	return -1;
}

static int psd_conn_udt_settimeout(psdConn *conn, uint8_t direction, int seconds) {
	psd_warn(10, "can't set timeout on udt sockets");
	return 0;
}

static void psd_conn_udt_free(psdConn *conn) {
	psdConn_udtData *udt_data = (psdConn_udtData *) conn->conn_private;

	free(udt_data);
}
